import torch
import torch.nn as nn
import torch.nn.functional as F


class IDCNN(nn.Module):

    eps = 1e-5

    def __init__(self, in_channels=1):
        super().__init__()

        # in_channels = 1

        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=3, padding=1, bias=False),
            nn.ReLU(inplace=True)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.conv5 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.conv6 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.conv7 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.conv8 = nn.Sequential(
            nn.Conv2d(64, in_channels, kernel_size=3, padding=1, bias=False),
            nn.ReLU(inplace=True)
        )
        self.last = nn.Tanh()

    def forward(self, x):
        input = self.conv1(x)
        input = self.conv2(input)
        input = self.conv3(input)
        input = self.conv4(input)
        input = self.conv5(input)
        input = self.conv6(input)
        input = self.conv7(input)
        noise = self.conv8(input)

        output = x / (noise + self.eps)

        output = self.last(output)

        return output, noise


class SAR_DRN(nn.Module):

    def __init__(self, in_channels=1):
        super().__init__()

        # in_channels = 1

        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels, 64, kernel_size=3, padding=1, 
                dilation=1, bias=False),
            nn.ReLU(inplace=True)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=2, 
                dilation=2, bias=False),
            nn.ReLU(inplace=True)
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=3, 
                dilation=3, bias=False),
            nn.ReLU(inplace=True)
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=4, 
                dilation=4, bias=False),
            nn.ReLU(inplace=True)
        )
        self.conv5 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=3, 
                dilation=3, bias=False),
            nn.ReLU(inplace=True)
        )
        self.conv6 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=2, 
                dilation=2, bias=False),
            nn.ReLU(inplace=True)
        )
        self.conv7 = nn.Sequential(
            nn.Conv2d(64, in_channels, kernel_size=3, padding=1, 
                dilation=1, bias=False)
        )

    def forward(self, x):
        input1 = self.conv1(x)
        input = self.conv2(input1)
        input2 = self.conv3(input)
        input2 = input1 + input2
        input3 = self.conv4(input2)
        input = self.conv5(input3)
        input4 = self.conv6(input)
        input4 = input3 + input4
        noise = self.conv7(input4)

        output = x - noise

        return output, noise


class RDNet(nn.Module):

    def __init__(self, in_channels=1):
        super().__init__()

        # in_channels = 1

        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=3, padding=1, bias=False),
            nn.ReLU(inplace=True)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.conv5 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        # self.conv6 = nn.Sequential(
        #     nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
        #     nn.BatchNorm2d(64),
        #     nn.ReLU(inplace=True)
        # )
        self.conv6 = nn.Sequential(
            nn.Conv2d(64, in_channels, kernel_size=3, padding=1, bias=False)
        )

    def forward(self, x):
        input = self.conv1(x)
        input = self.conv2(input)
        input = self.conv3(input)
        input = self.conv4(input)
        input = self.conv5(input)
        # input = self.conv6(input)
        noise = self.conv6(input)

        output = x - noise

        return output


if __name__ == "__main__":
    model = IDCNN()
    print(model)
    # model = SAR_DRN()
    # print(model)
    # model = RDNet()
    # print(model)
